
//封装的信号量对象，方便多线程情况下线程同步
class Semaphore {
    public:
     explicit Semaphore(int count = 0) : count_(count) {}

     void Signal() {
        std::unique_lock<std::mutex> lock(mutex_);
        ++count_;
        cv_.notify_one();
     }

     void Wait() {
        std::unique_lock<std::mutex> lock(mutex_);
        cv_.wait(lock, [=] { return count_ > 0; });
        --count_;
     }

    private:
     std::mutex mutex_;
     std::condition_variable cv_;
     int count_;
};

class TtsModelFact{
    public:
        TtsModelFact(std::string& model_path,std::string& model_fix);
        ~TtsModelFact();

        // Load model
        std::shared_ptr<tflite::FlatBufferModel> txtenc_model = nullptr;
        std::shared_ptr<tflite::FlatBufferModel> emot_model = nullptr;
        std::shared_ptr<tflite::FlatBufferModel> varadp_model = nullptr;
        std::shared_ptr<tflite::FlatBufferModel> duration_model = nullptr;
        std::shared_ptr<tflite::FlatBufferModel> LR_model = nullptr;
        std::shared_ptr<tflite::FlatBufferModel> dursin_model = nullptr;
        std::shared_ptr<tflite::FlatBufferModel> mem_model = nullptr;
        std::shared_ptr<tflite::FlatBufferModel> attenH_model = nullptr;
        std::shared_ptr<tflite::FlatBufferModel> getmask_model = nullptr;
        std::shared_ptr<tflite::FlatBufferModel> dec1_model = nullptr;
        std::vector<std::shared_ptr<tflite::FlatBufferModel>> dec2_pncas_model;
        std::shared_ptr<tflite::FlatBufferModel> dec2_model = nullptr;
        std::shared_ptr<tflite::FlatBufferModel> post_model = nullptr;
        std::shared_ptr<tflite::FlatBufferModel> post0_model = nullptr;
        std::shared_ptr<tflite::FlatBufferModel> post1_model = nullptr;
        std::shared_ptr<tflite::FlatBufferModel> post2_model = nullptr;
        std::shared_ptr<tflite::FlatBufferModel> post3_model = nullptr;
        std::shared_ptr<tflite::FlatBufferModel> dnorm_model = nullptr;

        std::shared_ptr<tflite::FlatBufferModel> hifigan0_model = nullptr;
        std::shared_ptr<tflite::FlatBufferModel> hifigan1_model = nullptr;
        std::shared_ptr<tflite::FlatBufferModel> hifigan2_model = nullptr;
        std::shared_ptr<tflite::FlatBufferModel> hifigan3_model = nullptr;
        std::shared_ptr<tflite::FlatBufferModel> hifigan4_model = nullptr;
        std::shared_ptr<tflite::FlatBufferModel> hifigan5_model = nullptr;

        std::shared_ptr<tflite::ops::builtin::BuiltinOpResolver> resolver = nullptr;

        std::vector<float> spk_embedding;
        int _pnca_num = 12;
};

class TtsModel{
    public:
        TtsModel(std::string& model_path,std::string &model_fix);
        ~TtsModel();

        // Load model
        std::shared_ptr<TtsModelFact> fact = nullptr;

        // each build can only create signle interpreter
        std::shared_ptr<tflite::InterpreterBuilder> txtenc_builder = nullptr;
        std::shared_ptr<tflite::InterpreterBuilder> emot_builder = nullptr;
        std::shared_ptr<tflite::InterpreterBuilder> varadp_builder = nullptr;
        std::shared_ptr<tflite::InterpreterBuilder> duration_builder = nullptr;
        std::shared_ptr<tflite::InterpreterBuilder> LR_builder = nullptr;
        std::shared_ptr<tflite::InterpreterBuilder> dursin_builder = nullptr;
        std::shared_ptr<tflite::InterpreterBuilder> mem_builder = nullptr;
        std::shared_ptr<tflite::InterpreterBuilder> attenH_builder = nullptr;
        std::shared_ptr<tflite::InterpreterBuilder> getmask_builder = nullptr;
        std::shared_ptr<tflite::InterpreterBuilder> dec1_builder = nullptr;
        std::vector<std::shared_ptr<tflite::InterpreterBuilder>> dec2_pncas_builder;
        std::shared_ptr<tflite::InterpreterBuilder> dec2_builder = nullptr;
        std::shared_ptr<tflite::InterpreterBuilder> post_builder = nullptr;
        std::shared_ptr<tflite::InterpreterBuilder> post0_builder = nullptr;
        std::shared_ptr<tflite::InterpreterBuilder> post1_builder = nullptr;
        std::shared_ptr<tflite::InterpreterBuilder> post2_builder = nullptr;
        std::shared_ptr<tflite::InterpreterBuilder> post3_builder = nullptr;
        std::shared_ptr<tflite::InterpreterBuilder> dnorm_builder = nullptr;
        std::shared_ptr<tflite::InterpreterBuilder> hifigan0_builder = nullptr;
        std::shared_ptr<tflite::InterpreterBuilder> hifigan1_builder = nullptr;
        std::shared_ptr<tflite::InterpreterBuilder> hifigan2_builder = nullptr;
        std::shared_ptr<tflite::InterpreterBuilder> hifigan3_builder = nullptr;
        std::shared_ptr<tflite::InterpreterBuilder> hifigan4_builder = nullptr;
        std::shared_ptr<tflite::InterpreterBuilder> hifigan5_builder = nullptr;
};

class stTTS{
    public:
        stTTS(std::string& model_path,std::string& model_fix, void (*pcm_callback)(std::vector<short> &pcm,void* args,bool endf),void *args);
        ~stTTS();
        void input_txt(const std::vector<std::vector<std::vector<std::string>>> &sentences);
    public:
        std::shared_ptr<TtsModel> _model = nullptr;
        std::unique_ptr<tflite::Interpreter> txtenc_interpreter;
        std::unique_ptr<tflite::Interpreter> emot_interpreter;
        std::unique_ptr<tflite::Interpreter> varadp_interpreter;
        std::unique_ptr<tflite::Interpreter> duration_interpreter;
        std::unique_ptr<tflite::Interpreter> LRtxt_interpreter;
        std::unique_ptr<tflite::Interpreter> LRspk_interpreter;
        std::unique_ptr<tflite::Interpreter> LRemo_interpreter;
        std::unique_ptr<tflite::Interpreter> dursin_interpreter;
        std::unique_ptr<tflite::Interpreter> mem_interpreter;
        std::unique_ptr<tflite::Interpreter> attenH_interpreter;
        std::unique_ptr<tflite::Interpreter> getmask_interpreter;
        std::unique_ptr<tflite::Interpreter> dec1_interpreter;
        std::vector<std::unique_ptr<tflite::Interpreter>> dec2_pncas_interpreter;
        std::unique_ptr<tflite::Interpreter> dec2_interpreter;
        std::unique_ptr<tflite::Interpreter> post_interpreter;
        std::unique_ptr<tflite::Interpreter> post0_interpreter;
        std::unique_ptr<tflite::Interpreter> post1_interpreter;
        std::unique_ptr<tflite::Interpreter> post2_interpreter;
        std::unique_ptr<tflite::Interpreter> post3_interpreter;
        std::unique_ptr<tflite::Interpreter> dnorm_interpreter;
        std::unique_ptr<tflite::Interpreter> hifigan0_interpreter;
        std::unique_ptr<tflite::Interpreter> hifigan1_interpreter;
        std::unique_ptr<tflite::Interpreter> hifigan2_interpreter;
        std::unique_ptr<tflite::Interpreter> hifigan3_interpreter;
        std::unique_ptr<tflite::Interpreter> hifigan4_interpreter;
        std::unique_ptr<tflite::Interpreter> hifigan5_interpreter;
    private:
        void init();
        void clear();
        bool get_sent_ids(const std::vector<std::vector<std::string>> &sent, std::vector<std::vector<int>> &sentid);
        void am_dec(float *memdata, int memlen,int memdim, int durmax);
        bool am_post();
        static void worker_vocoder(stTTS *);
        void st_vocoder1(bool endflag);
        void st_vocoder2(bool _end);
        void st_vocoder3(bool _end);
        void st_vocoder4(bool _end);
        void st_vocoder5(bool _end);
        int _meldim;
        int _postdim;
        int _postlen = 0;
        int _hifi0len = 0;
        int _hifi1len = 0;
        int _hifi3len = 0;
        int _hifi4len = 0;
        bool _endflag;
        int _pad_memlen = 0;
        int _post0_idx = 0;
        int _post_idx = 0;
        std::vector<float> _post_ph;
        std::vector<float> _post_pc;
        std::vector<float> _dec2output;
        std::vector<float> _post0_lpad;
        std::vector<float> _post1_lpad;
        std::vector<float> _post2_lpad;
        std::vector<float> _post3_lpad;
        std::vector<float> _post0out;
        std::vector<float> _post1out;
        std::vector<float> _post2out;
        std::vector<float> _post3out;
        std::vector<float> _dnormout;
        std::vector<float> _hifi0out;
        std::vector<float> _exciout;
        int _excidim=1;
        int _hifi0olp = 3;
        int _hifi0win = 24;
        int _hifi1olp = 3;
        int _hifi1idx = _hifi1olp;
        int _hifi1dim = 256;
        int _hifi1win = 21;
        int _hifi1up = 10;
        std::vector<float> _hifi1out;
        int _hifi2olp = 5;
        int _hifi2idx = _hifi2olp;
        int _hifi2dim = 128;
        int _hifi2win = 150;
        int _hifi2up = 5;
        std::vector<float> _hifi2out;
        int _hifi3olp = 20;
        int _hifi3idx = _hifi3olp;
        int _hifi3dim = 64;
        int _hifi3win = 700;
        int _hifi3up = 2;
        std::vector<float> _hifi3out;
        int _hifi4olp = 20;
        int _hifi4idx = _hifi4olp;
        int _hifi4dim = 32;
        int _hifi4win = 1320;
        int _hifi4up = 2;
        std::vector<float> _hifi4out;
        int _hifi5dim = 16;
        int _hifi5olp = 4;
        std::vector<short> _pcm;
        int _pcmlen = 0;
        int _pcmid=0;

        Semaphore _vocsmp;
        std::mutex _dec_mutex;
        std::shared_ptr<std::thread> th_voc = nullptr;
        void (*_pcm_callback)(std::vector<short> &pcm,void* args,bool endf);
        void *_args;
        std::chrono::time_point<std::chrono::steady_clock> _start;
};
